from __future__ import with_statement

import random
import transaction
import time
from crv.model import DBSession, metadata
from crv.model.certstore_db import *
from crv.tools.singleton import Singleton
from M2Crypto import X509, EVP, RSA, Rand, ASN1, m2, util, BIO
import os, time, base64, sys
import uuid

import logging
log = logging.getLogger(__name__)

class CertStoreMgr(Singleton):
    """
    @author Yushu Yao
    This CertStoreMgr will only manage one set of Certs. At __init__ it will look into its storage backend, and find the first CertStore and use it.
    The certstore with a unique CA name is created when running paster setup-app
    The CA name is CloudCRV_CA+UUID
    
    Note: We also assume that CertStoreMgr is the only one that's working on the same storagebackend
    """

    def __initmembers(self):
        self.__caname=None
        self.__capassword = None
        self.__carequest = None
        self.__cacert = None

        #M2crypto.RSA type
        self.__caprivatekey =None

        #public key will be generated from private key on the fly
#        self.__capublickey =None
        self.__certs=None

    def __init__(self):
        super(CertStoreMgr, self).__init__()
        
        self.__initialized=False
        self.__initmembers()

        #use the database backend
        self.__storage=DB_CertStoreStorage()

        #look into storage and find the first CertStore
        certstorenames=self.__storage.getCAList()

        if len(certstorenames)>=1:
            self.__caname=certstorenames[0]
            #load the info for the ca and cache it for quick access
            #note the ca cert has a cn=caname

            self.__capassword = self.__storage.loadCAPassword(self.__caname)
            (strcarequest,strcacert, strcaprivatekey, strpassword) = self.__storage.loadCert(self.__caname, self.__caname)

#            print "self.__capassword="+self.__capassword+" strpassword="+strpassword
    
            self.__carequest = X509.load_request_string(strcarequest)
            self.__cacert = X509.load_cert_string(strcacert)

            self.__caprivatekey=RSA.load_key_string(strcaprivatekey, callback=self.__capasscallback)

            self.__certs={}

            #cache in all existing certs
            certnames=self.__storage.getCertList(self.__caname)

            for icertname in certnames:
                if icertname==self.__caname:
                    continue
                (strrequest,strcert, strprivatekey, strpassword) = self.__storage.loadCert(self.__caname, icertname)
                def passwordcallback(*args):
                    return strpassword
                prikey=RSA.load_key_string(strprivatekey, callback=passwordcallback)
                self.__certs[icertname]=[X509.load_request_string(strrequest), 
                                         X509.load_cert_string(strcert), 
                                         prikey,
                                         strpassword]
                
            self.__initialized=True

    def initializeCA(self, caname=None, password=None):
        """Used to generate a new ca cert"""
        if self.__initialized:
            raise RuntimeError("already initizlized, with caname=%s, try to remove storage backend(database) entries"%self.__caname)

        self.__caname=caname
        if caname is None:
            #randome a caname
            self.__caname="cloudcrv_ca_%s"%str(uuid.uuid4())
        
        if password is None:
            self.__capassword=str(uuid.uuid4()).replace("-","")[0:15]
        else:
            self.__capassword=password

        (self.__carequest,self.__cacert, caprivatekey_pkey, capublickey_pkey) = self.__mkcacert(self.__caname, self.__capassword)
        
        self.__caprivatekey=self.__PKEY2RSA(caprivatekey_pkey)

        self.__certs={}
        
        self.__storage.newCA(self.__caname, self.__capassword)
        self.__storage.saveCert(self.__caname, self.__caname, self.__carequest.as_pem(), self.__cacert.as_pem(), self.__caprivatekey.as_pem(callback=self.__capasscallback), self.__capassword, overwrite=True)

        self.__initialized=True

    def __RSA2PKEY(self, thersa):
        def callback(*arg):
            return "randompass"
        mem=BIO.MemoryBuffer()
        thersa.save_key_bio(mem,callback=callback)
        return EVP.load_key_bio(mem, callback=callback)

    def __PKEY2RSA(self, thepkey):
        def callback(*arg):
            return "randompass"
        mem=BIO.MemoryBuffer()
        thepkey.save_key_bio(mem,callback=callback)
        return RSA.load_key_bio(mem, callback=callback)

    def getCACertPEM(self):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        return self.__cacert.as_pem()

    def getCARequestPEM(self):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        return self.__carequest.as_pem()

    def getCAPrivateKeyPEM(self):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        return self.__caprivatekey.as_pem(callback=self.__capasscallback)

    def getCAPublicKeyPEM(self):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        mem=BIO.MemoryBuffer()
        self.__caprivatekey.save_pub_key_bio(mem)
        return mem.getvalue()

    def isInitialized(self):
        return self.__initialized

    def getCAName(self):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        return self.__caname

    def getCertList(self):
        return self.__certs.keys()

    def getCAPassword(self):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        return self.__capassword

    def destroyCA(self):
        """Remove all trace of this cert store with this caname"""
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        self.__storage.removeCertStore(self.__caname)
        self.__initmembers()
        
    def generateCert(self, cn):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")

        if cn in self.__certs.keys():
            raise RuntimeError("Cert with CN=%s already exist in CertStore CAName=%s"%(cn, self.__caname))

        #generate a password
        password=str(uuid.uuid4()).replace("-","")[0:15]
        
        log.debug("password=%s"%password)
        log.debug("cacert=%s"%self.__cacert.as_pem())
        log.debug("caprik=%s"%self.__caprivatekey.as_pem(callback=self.__capasscallback))

        #first create the cert
        (cert, privatekey, publickey, req)=self.__mkcert(self.__cacert, self.__RSA2PKEY(self.__caprivatekey), cn, password=password)
        
        self.__certs[cn]=[req, cert, self.__PKEY2RSA(privatekey), password]

        def passwordcallback(*args):
            return password

        #now save to storage backend
        self.__storage.saveCert(self.__caname, cn, req.as_pem(), cert.as_pem(), privatekey.as_pem(callback=passwordcallback), password=password, overwrite=True)

    def getPassword(self, cn):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        # format of cert: self.__certs[cn]=[req, cert, privatekey, password]
        if cn in self.__certs.keys():
            return self.__certs[cn][3]
        return None
        
    def getCertPEM(self, cn):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        # format of cert: self.__certs[cn]=[req, cert, privatekey, password]
        if cn in self.__certs.keys():
            return self.__certs[cn][1].as_pem()
        return None
    
    def getRequestPEM(self, cn):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        # format of cert: self.__certs[cn]=[req, cert, privatekey, password]
        if cn in self.__certs.keys():
            return self.__certs[cn][0].as_pem()
        return None
        
    def getPrivateKeyPEM(self, cn):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        # format of cert: self.__certs[cn]=[req, cert, privatekey, password]
        if cn in self.__certs.keys():
            def passwordcallback(*args):
                return self.__certs[cn][3]
            return self.__certs[cn][2].as_pem(callback=passwordcallback)
        return None
        
    def getPublicKeyPEM(self, cn):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        # format of cert: self.__certs[cn]=[req, cert, privatekey, password]
        if cn in self.__certs.keys():
            def passwordcallback(*args):
                return self.__certs[cn][3]

            mem=BIO.MemoryBuffer()
            self.__certs[cn][2].save_pub_key_bio(mem)
            return mem.getvalue()
        return None
    
    def revoke(self, cn):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        pass
    
    def delete(self, cn):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        if not cn in self.__certs.keys():
            raise RuntimeError("Cert CN=%s does not exist in Certstore caname=%s"%(cn, self.__caname))
        certentry=self.__certs.pop(cn)
        #remove from storage backend
        self.__storage.removeCert(self.__caname, cn)
        

    def verify(self, certpemstr):
        if not self.__initialized:
            raise RuntimeError("Not yet initialized")
        #cert as pem string
        cert=X509.load_cert_string(certpemstr)

        if not cert.get_subject().CN in self.__certs.keys():
            log.warning("verify::CN=%s not in the list of certs"%cert.get_subject().CN)
        
        if cert.verify(self.__cacert.get_pubkey()) == 1:
            return True

        return False

        
    def __capasscallback(self,*args):
        return self.__capassword
    
    def __mkreq(self, bits, namecn, password):
        pk = EVP.PKey()
        x = X509.Request()
        def passwordcallback(*args):
            return password
        rsa = RSA.gen_key(bits, 65537, callback=passwordcallback)
        pk.assign_rsa(rsa)
        rsa = None # should not be freed here
        x.set_pubkey(pk)
        x.set_subject(self.__base_dn(namecn))
        x.sign(pk,'sha1')
        pk2 = x.get_pubkey()
        return x, pk
    
    def __mkcacert(self,namecn, password):
        req, pk = self.__mkreq(2048, namecn=namecn, password=password)
        pkey = req.get_pubkey()
        sub = req.get_subject()
        cert = X509.X509()
        cert.set_serial_number(random.randint(1,9999999))
        cert.set_version(1)
        cert.set_subject(sub)
        t = long(time.time()) + time.timezone
        now = ASN1.ASN1_UTCTIME()
        now.set_time(t)
        nowPlusYear = ASN1.ASN1_UTCTIME()
        nowPlusYear.set_time(t + 60 * 60 * 24 * 365 * 5)
        cert.set_not_before(now)
        cert.set_not_after(nowPlusYear)
        self.__add_ext(cert)
        cert.set_issuer(self.__base_dn(namecn))
        cert.set_pubkey(pkey) 
        cert.sign(pk, 'sha1')
        return req, cert, pk, pkey
    
    def __add_ext(self, cert):
        ext2 = X509.new_extension('nsComment', 'For more information about CloudCRV, please refer to http://code.google.com/p/cloudcrv/, or contact author: yao.yushu@gmail.com')
        cert.add_ext(ext2)
    
    def __base_dn(self, CN):
        issuer = X509.X509_Name()
        issuer.C = "US"
        issuer.ST = "CA"
        issuer.L = "Berkeley"
        issuer.O = "Lawrence Berkeley National Laboratory"
        issuer.CN = CN
        return issuer
    
    def __mkcert(self,cacert, cakey, CN, password):
        req, pk = self.__mkreq(2048, namecn=CN, password=password)
        log.debug("cn=%s"%CN)

        pkey = req.get_pubkey()
    
        sub = req.get_subject()
    
        cert = X509.X509()
        cert.set_serial_number(random.randint(1,9999999))
        cert.set_version(2)
        cert.set_subject(sub)
    
        t = long(time.time()) + time.timezone
        now = ASN1.ASN1_UTCTIME()
        now.set_time(t)
        nowPlusYear = ASN1.ASN1_UTCTIME()
        nowPlusYear.set_time(t + 60 * 60 * 24 * 365 * 5)
    
        cert.set_not_before(now)
        cert.set_not_after(nowPlusYear)
    
        issuer = cacert.get_subject()
    
        cert.set_issuer(cacert.get_subject())
        cert.set_pubkey(pkey)
        self.__add_ext(cert)
    
        log.debug("capass=%s, cakey=%s"%(self.__capasscallback(),cakey.as_pem(callback=self.__capasscallback)))
#        try:
        cert.sign(cakey, 'sha1')
#        except AttributeError:
#            pass
        
        return cert,pk,pkey,req
    
        
class ICertStoreStorage:
    
    def __init__(self):
        pass

    def saveDoc(self, caname, doc):
        raise NotImplementedError

    def loadDoc(self, caname):
        raise NotImplementedError

    def saveMisc(self, caname, misc):
        raise NotImplementedError

    def loadMisc(self, caname):
        raise NotImplementedError

    def saveCAPassword(self, caname, password):
        raise NotImplementedError

    def loadCAPassword(self, caname):
        raise NotImplementedError

    def saveCert(self,caname, cn, request, certificate, privatekey, overwrite=False):
        raise NotImplementedError

    def loadCert(self, caname, cn):
        raise NotImplementedError

    def existCert(self, caname, cn):
        raise NotImplementedError

    def existCA(self, caname):
        raise NotImplementedError

    def getCAList(self):
        raise NotImplementedError

    def removeCert(self, caname, cn):
        raise NotImplementedError

    def removeCertStore(self, caname):
        raise NotImplementedError

    def getCertList(self, caname):
        raise NotImplementedError

class DB_CertStoreStorage(ICertStoreStorage):
    
    def __init__(self):
        pass

    def saveDoc(self, caname, doc):
        cas = DBSession.query(CertStore).filter(CertStore.caname==caname).all()
        if len(cas)<1:
            raise RuntimeError("saveDoc: caname=%s doesn't exist"%caname)
        cas[0].doc=doc
        #transaction.commit()

    def loadDoc(self, caname):
        cas = DBSession.query(CertStore).filter(CertStore.caname==caname).all()
        if len(cas)<1:
            raise RuntimeError("loadDoc: caname=%s doesn't exist"%caname)
        return cas[0].doc

    def saveMisc(self, caname, misc):
        cas = DBSession.query(CertStore).filter(CertStore.caname==caname).all()
        if len(cas)<1:
            raise RuntimeError("saveMisc: caname=%s doesn't exist"%caname)
        cas[0].misc=misc
        #transaction.commit()

    def loadMisc(self, caname):
        cas = DBSession.query(CertStore).filter(CertStore.caname==caname).all()
        if len(cas)<1:
            raise RuntimeError("loadMisc: caname=%s doesn't exist"%caname)
        return cas[0].misc

    def saveCAPassword(self, caname, password):
        cas = DBSession.query(CertStore).filter(CertStore.caname==caname).all()
        if len(cas)<1:
            raise RuntimeError("saveCAPassword: caname=%s doesn't exist"%caname)
        cas[0].password=password
        #transaction.commit()

    def loadCAPassword(self, caname):
        cas = DBSession.query(CertStore).filter(CertStore.caname==caname).all()
        if len(cas)<1:
            raise RuntimeError("loadCAPassword: caname=%s doesn't exist"%caname)
        return cas[0].password

    def saveCert(self,caname, cn, request, certificate, privatekey, password, overwrite=False):
        
        #first figure out if caname, cn already exist
        certs = DBSession.query(Certificate).join(CertStore).filter(CertStore.caname==caname).filter(Certificate.cn==cn).all()

        if len(certs)>=1:
            if not overwrite:
                raise RuntimeError("saveCert: caname=%s cn=%s already exist, use overwrite=True to overwrite"%(caname, cn))

            certs[0].cn=cn
            certs[0].request=request
            certs[0].certificate=certificate
            certs[0].privatekey=privatekey
            certs[0].password=password

            #transaction.commit()

        else:
            certstores = DBSession.query(CertStore).filter(CertStore.caname==caname).all()
            if len(certstores)<1:
                raise RuntimeError("saveCert: caname=%s doesn't exist"%caname)

            thecert=Certificate()
            thecert.cn=cn
            thecert.request=request
            thecert.certificate=certificate
            thecert.privatekey=privatekey
            thecert.certstore=certstores[0]
            thecert.password=password
            
            #transaction.commit()

    def loadCert(self, caname, cn):
        certs = DBSession.query(Certificate).join(CertStore).filter(CertStore.caname==caname).filter(Certificate.cn==cn).all()
        if len(certs)<1:
            raise RuntimeError("loadCert: caname=%s cn=%s does not exist"%(caname, cn))
        return (certs[0].request, certs[0].certificate, certs[0].privatekey, certs[0].password)
        
    def existCert(self, caname, cn):
        ncerts = DBSession.query(Certificate).join(CertStore).filter(CertStore.caname==caname).filter(Certificate.cn==cn).count()
        if ncerts<1:
            return False
        return True

    def removeCert(self, caname, cn):
        certs = DBSession.query(Certificate).join(CertStore).filter(CertStore.caname==caname).filter(Certificate.cn==cn).all()
        if len(certs)<1:
            raise RuntimeError("removeCert: caname=%s cn=%s does not exist"%(caname, cn))
        for icerts in certs:
            DBSession.delete(icerts)
        #transaction.commit()

    def removeCertStore(self, caname):
        certs = DBSession.query(Certificate).join(CertStore).filter(CertStore.caname==caname).all()
        cas=DBSession.query(CertStore).filter(CertStore.caname==caname).all()
        if len(cas)<1:
            raise RuntimeError("removeCertStore: caname=%s does not exist"%(caname))
        for icerts in certs:
            DBSession.delete(icerts)
        DBSession.delete(cas[0])
        #transaction.commit()

    def existCA(self, caname):
        ncas = DBSession.query(CertStore).filter(CertStore.caname==caname).count()
        if ncas<1:
            return False
        return True

    def newCA(self, caname, password, doc="", misc=""):
#        print "NewCA %s %s %s %s"%(caname, password, doc, misc)
        certstore=CertStore()
        certstore.caname=caname
        certstore.password=password
        certstore.doc=doc
        certstore.misc=misc
        DBSession.add(certstore)
        #transaction.commit()

    def getCAList(self):
        """Return a list of strings with canames inside this storage """
        cas = DBSession.query(CertStore.caname).all()
        if len(cas)>=1:
            return map(lambda x: str(x[0]), cas) 
        return []

    def getCertList(self, caname):
        certs = DBSession.query(Certificate.cn).join(CertStore).filter(CertStore.caname==caname).all()
        if len(certs)>=1:
            return map(lambda x: str(x[0]), certs) 
        return []

    
